from pymysql import connect, Connect
from pymysql.cursors import DictCursor, Cursor
from typing import Dict, List
import pandas as pd


def get_db_connection(
    host="localhost", port=3306, user="root", password=None, database="test"
):
    conn = connect(
        host=host, user=user, port=port, password=password, database=database
    )
    return conn


def insert_to_db(conn: Connect, data: dict, table_name: str) -> int | None:
    """
    Args:
        conn: pymysql的连接实例
        data: 待插入表的python dict
        tablename: 表名
    """
    cursor = conn.cursor()
    columns = ", ".join(data.keys())
    placeholders = ", ".join(["%s"] * len(data))
    sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
    values = tuple(data.values())
    cursor.execute(sql, values)
    conn.commit()
    return cursor.lastrowid


def search(conn: Connect, table: str, where) -> List[Dict]:
    sql = "select * from %s where %s" % (table, where)
    cursor = conn.cursor(DictCursor)
    try:
        cursor.execute(sql)
        ret = cursor.fetchall()
    except:
        ret = []
    return ret


def to_excel(conn: Connect, sql: str, filename: str, index=False):
    """
    @param {pymysql.Connect} conn - 已经建立的mysql连接实例
    @param {str} sql - 用于从mysql中检索出数据的SQL语句
    @param {str} filename - 需要保存的excel的文件名
    @returns {None} 成功没有输出
    @index {bool} [index] - 是否保存行号, 默认值：False
    @example
    to_excel(conn, 'select * from library', 'library.xlsx', index=True)
    """
    df = pd.read_sql(sql, con=conn)
    df.to_excel(filename, index=index)
